import os
import tensorflow as tf
from PIL import Image
 
# 图像分类
def class_text_to_int(row_label):
    if row_label == 'negative':
        return 2
    elif row_label == 'positive':
        return 1
 
# 创建TFRecord
def convert_to_TFRecord(input_dir, output_file, image_format='.jpg'):
    writer = tf.io.TFRecordWriter(output_file)
    for sub_dir in ['negative', 'positive']:
        class_id = class_text_to_int(sub_dir)
        dir_path = os.path.join(input_dir, sub_dir)
        for image_name in os.listdir(dir_path):
            image_path = os.path.join(dir_path, image_name)
            img = Image.open(image_path)
            img = img.resize((227, 227))
            img_raw = img.tobytes()
            example = tf.train.Example(features=tf.train.Features(feature={
                'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[class_id])),
                'image_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
            }))
            writer.write(example.SerializeToString())

    writer.close()

if __name__ == '__main__':
    # 调用函数，转换数据集
    input_dir = 'E:/workspace/web/www/zzd/daelui-tensorflow/crack_detection/dataset/Concrete Crack Images for Classification/train'  # 你的数据集路径
    output_file = 'E:/workspace/web/www/zzd/daelui-tensorflow/crack_detection/tf_record/train.tfrecord'  # 输出的TFRecord文件名
    convert_to_TFRecord(input_dir, output_file)